// Copyright (c) 2019 Presto Labs Pte. Ltd.
// Author: jaewon

#include "coin2/base/math/twma.h"

#include <memory>

#include <gtest/gtest.h>

namespace coin2::base::math {

TEST(TimeWeightedMovingAverageTest, EasyTest) {
  TimeWeightedMovingAverage twma(5);

  EXPECT_EQ(twma.tw_sum(), 0.);

  twma.Push(0, 1.);
  EXPECT_EQ(twma.tw_sum(), 0.);

  twma.PushTime(1);
  EXPECT_EQ(twma.tw_sum(), 1.);

  twma.PushTime(3);
  EXPECT_EQ(twma.tw_sum(), 3.);

  twma.PushTime(5);
  EXPECT_EQ(twma.tw_sum(), 5.);

  twma.PushTime(7);
  EXPECT_EQ(twma.tw_sum(), 5.);

  twma.PushTime(10);
  EXPECT_EQ(twma.tw_sum(), 5.);

  twma.Push(15, 2.);
  EXPECT_EQ(twma.tw_sum(), 5.);

  twma.Push(17, 3.);
  EXPECT_EQ(twma.tw_sum(), 7.);

  twma.Push(18, 4.);
  EXPECT_EQ(twma.tw_sum(), 9.);

  twma.Push(19, 5.);
  EXPECT_EQ(twma.tw_sum(), 12.);

  twma.PushTime(21);
  EXPECT_EQ(twma.tw_sum(), 19.);

  twma.PushTime(22);
  EXPECT_EQ(twma.tw_sum(), 22.);

  twma.PushTime(23);
  EXPECT_EQ(twma.tw_sum(), 24.);

  twma.PushTime(24);
  EXPECT_EQ(twma.tw_sum(), 25.);

  twma.PushTime(25);
  EXPECT_EQ(twma.tw_sum(), 25.);

  twma.PushTime(100);
  EXPECT_EQ(twma.tw_sum(), 25.);
}

void AverageTest() {
  const int test_size = 1000000;
  const int window = 17293;
  std::unique_ptr<double[]> data(new double[test_size]);
  std::unique_ptr<double[]> cum(new double[test_size]);

  data[0] = 0;
  cum[0] = 0;
  for (int i = 1; i < test_size; i++) {
    int64_t elem = i % 85293;
    data[i] = (elem * elem * elem) % 52364;
    cum[i] = cum[i - 1] + data[i - 1];
  }

  TimeWeightedMovingAverage twma(window);
  double last_value = -1.;
  for (int i = 0; i < test_size; i++) {
    if (last_value != data[i]) {
      twma.Push(i, data[i]);
      last_value = data[i];
    } else {
      twma.PushTime(i);
    }

    if (i >= window) {
      ASSERT_EQ(twma.tw_sum(), cum[i] - cum[i - window]);
    } else {
      ASSERT_EQ(twma.tw_sum(), cum[i]);
    }
  }
}

TEST(TimeWeightedMovingAverageTest, AverageTest) { AverageTest(); }

}  // namespace coin2::base::math
